#include "HixlCsClient.h"
#include "HixlMemStore.h"
#include "HixlEventLoop.h"
#include <cstring>
#include <vector>

namespace hixl {
    // ============================================================
    // Constructor / Destructor
    // ============================================================

    // HixlCsClient：构造函数
    // ------------------------------------------------------------
    // 初始化 EndpointStore，并重置连接状态。
    // 一个 client 对象对应一条可能的 socket + endpoint + channel 流程。
    // ------------------------------------------------------------
    HixlCsClient::HixlCsClient()
        : endpoint_store_(std::make_unique<EndpointStore>()) {
        ResetConnectState();
    }

    // ~HixlCsClient：析构
    // ------------------------------------------------------------
    // 析构时确保释放所有 endpoint / channel / mem 导入信息。
    // Destroy() 主动销毁；析构只是兜底。
    // ------------------------------------------------------------
    HixlCsClient::~HixlCsClient() {
        DestroyAllEndpointsAndChannels();
        ClearRemoteMemInfo();
    }


    // ============================================================
    // Create
    // ============================================================

    // Create：创建 Client 初始状态
    // ------------------------------------------------------------
    // 负责：
    // 1）记录 server_ip/server_port
    // 2）保存本端 Endpoint（self_endpoint_）
    // 3）在 EndpointStore 中创建本端 EndpointHandle
    // Create 并不发起 connect，仅完成“静态初始化”。
    // ------------------------------------------------------------
    hccl::HcclResult HixlCsClient::Create(char *server_ip,
                                          uint32_t server_port,
                                          const hccl::EndPoint *end_point) {
        if (server_ip == nullptr || end_point == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        std::lock_guard<std::mutex> lock(mutex_);

        // 初始化 Client 内部状态
        auto ret = InitState(server_ip, server_port, end_point);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }

        // 创建本端 Endpoint 对象（self_endpoint_handle_）
        return CreateSelfEndpoint(end_point);
    }


    // ============================================================
    // Connect
    // ============================================================

    // Connect：发起非阻塞建链
    // ------------------------------------------------------------
    // 负责：
    // 1）第一次调用 Connect() 时，创建 socket（InitSocketIfNeeded）
    // 2）调用 HcclSocket::Connect() 发起非阻塞建链
    // 3）更新状态为 RA_CONNECTING
    // Connect 不会等待建链完成，只是“开始”。
    // ------------------------------------------------------------
    hccl::HcclResult HixlCsClient::Connect() {
        std::lock_guard<std::mutex> lock(mutex_);

        // 已建链完成
        if (phase_ == ConnectPhase::DONE && is_connected_) {
            return hccl::HcclResult::SUCCESS;
        }
        // 已失败
        if (phase_ == ConnectPhase::FAILED) {
            return hccl::HcclResult::E_PARA;
        }

        // 第一次调用 Connect()
        if (!connect_started_) {
            // 创建 socket
            auto ret = InitSocketIfNeeded();
            if (ret != hccl::HcclResult::SUCCESS) {
                return ret;
            }

            // 发起非阻塞 connect
            ret = StartRaConnect();
            if (ret != hccl::HcclResult::SUCCESS) {
                phase_ = ConnectPhase::FAILED;
                return ret;
            }

            connect_started_ = true;
            phase_ = ConnectPhase::RA_CONNECTING;
        }

        return hccl::HcclResult::SUCCESS;
    }


    // ============================================================
    // GetStatus
    // ============================================================

    // GetStatus：查询当前连接状态
    // ------------------------------------------------------------
    // 用户反复调用该函数（非阻塞轮询），用来推动状态机：
    // Connect() ——> RA_CONNECTING ——> HANDSHAKE_SEND_EP ——> HANDSHAKE_RECV_EP ——>
    // CHANNEL_CREATING ——> DONE
    //
    // 这里主要根据 socket_->GetStatus() 决定下一步。
    // ------------------------------------------------------------
    hccl::HcclResult HixlCsClient::GetStatus(int32_t *status) {
        std::lock_guard<std::mutex> lock(mutex_);
        if (!connect_started_) {
            *status = 0;
            return hccl::HcclResult::SUCCESS;
        }

        // 处理“connect 正在建立中”
        if (phase_ == ConnectPhase::RA_CONNECTING && socket_) {
            auto ret = HandleRaConnecting(status);
            if (ret != hccl::HcclResult::SUCCESS) {
                return ret;
            }
        } else {
            // 处理握手 / channel 建立阶段
            UpdateStatusByPhase(status);
        }

        return hccl::HcclResult::SUCCESS;
    }


    // ============================================================
    // GetRemoteMem
    // ============================================================

    // GetRemoteMem：获取服务端共享内存信息
    // ------------------------------------------------------------
    // 流程：
    // 1）确保连接已完成（is_connected_）
    // 2）找到当前活跃的 socket + endpoint
    // 3）SendRemoteMemRequest（请求 server 的 memdesc）
    // 4）RecvRemoteMemCount（获得共享内存个数）
    // 5）循环接收每个 memdesc + tag，并执行 MemImport
    // ------------------------------------------------------------
    hccl::HcclResult HixlCsClient::GetRemoteMem(hccl::HcclMem **remote_mem_list,
                                                char **mem_tag_list,
                                                uint32_t *list_num,
                                                uint32_t timeout) {
        if (remote_mem_list == nullptr ||
            mem_tag_list == nullptr ||
            list_num == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        if (!is_connected_) {
            return hccl::HcclResult::E_PARA;
        }

        // 获取活跃的 {socket, endpoint}
        std::shared_ptr<hccl::HcclSocket> socket;
        EndpointHandle endpoint = nullptr;
        auto ret = GetActiveSocketAndEndpoint(socket, endpoint);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }

        // 发送请求
        ret = SendRemoteMemRequest(*socket, timeout);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }

        // 接收内存个数 + 分配数组
        uint32_t memCount = 0;
        ret = RecvRemoteMemCount(*socket, &memCount,
                                 remote_mem_list, mem_tag_list);
        if (ret != hccl::HcclResult::SUCCESS || memCount == 0) {
            *list_num = memCount;
            return ret;
        }

        // 循环接收每条共享内存描述
        hccl::HcclMem *memList = *remote_mem_list;
        char *tagBuf = *mem_tag_list;

        for (uint32_t i = 0; i < memCount; ++i) {
            char *slot = tagBuf + i * TAG_MAX_LEN;
            ret = RecvOneMemEntry(*socket, endpoint, memList[i],
                                  slot, TAG_MAX_LEN);
            if (ret != hccl::HcclResult::SUCCESS) {
                CleanupRemoteMemOnError(memList, tagBuf);
                *remote_mem_list = nullptr;
                *mem_tag_list = nullptr;
                *list_num = 0;
                return ret;
            }
        }

        *list_num = memCount;
        return hccl::HcclResult::SUCCESS;
    }

    // 注册client的endpoint的内存信息到内存注册表中。                              
    hccl::HcclResult ClientRegMem(const char *mem_tag, const HcclMem *mem, MemHandle *mem_handle) {

    }

    //通过已经建立好的channel，从用户提取的地址列表中，批量读取server内存地址中的内容
    hccl::HcclResult BatchGet(uint32_t list_num, void** remote_buf_list, 
                        const void** local_buf_list, uint64_t* len_list) {

                        }    
            
    //通过已经建立好的channel，从用户提取的地址列表中，批量往server内存地址中写入内容
    hccl::HcclResult BatchPut(uint32_t list_num, void** remote_buf_list, 
                        const void** local_buf_list, uint64_t* len_list) {

                        }
                        
    //通过已经建立好的channel，检查批量读写的状态。
    hccl::HcclResult CheckStatus() {

    }                     

    // 注册client的endpoint的内存信息到内存注册表中。                              
    hccl::HcclResult ClientUnRegMem() {

    }

    // ============================================================
    // Destroy
    // ============================================================

    // Destroy：销毁 client 实例
    // ------------------------------------------------------------
    // 负责：
    // 1）从 epoll 中删除客户端 socket
    // 2）向 server 发送“我要断开”的控制消息
    // 3）关闭 socket + 删除 endpoints/channels
    // 4）清理 RemoteMemInfo
    // 5）重置内部状态
    //
    // Destroy 后，这个 client 相当于彻底不可再用。
    // ------------------------------------------------------------

    hccl::HcclResult HixlCsClient::Destroy() {
        std::lock_guard<std::mutex> lock(mutex_);

        // 1）从 epoll 中移除
        if (registered_fd_ != nullptr) {
            HixlEventLoop::GetInstance().DelClient(registered_fd_);
            registered_fd_ = nullptr;
        }

        // 2）通知 server 我们要断开
        if (socket_) {
            static const char closeMsg[] = "HIXL_CS_CLIENT_CLOSE";
            (void) socket_->Send(closeMsg, sizeof(closeMsg));

            // 设置为强制关闭
            socket_->SetForceClose(true);
        }

        // 3）销毁所有 endpoint + channel + socket
        DestroyAllEndpointsAndChannels();

        // 4）清理 MemImport 记录
        ClearRemoteMemInfo();

        // 5）重置状态（相当于销毁 clientHandle）
        ResetConnectState();

        return hccl::HcclResult::SUCCESS;
    }


    // ============================================================
    // OnSocketEvent
    // ============================================================

    // OnSocketEvent：epoll 事件回调
    // ------------------------------------------------------------
    // 当 epoll 事件触发（可写 / 可读）时，EventLoop 会调用此函数。
    // 我们根据当前 phase_ 决定推进握手或 channel 建立。
    // ------------------------------------------------------------
    hccl::HcclResult HixlCsClient::OnSocketEvent(hccl::FdHandle fd) {
        std::lock_guard<std::mutex> lock(mutex_);

        if (!socket_ || socket_->GetFdHandle() != fd) {
            return hccl::HcclResult::E_PARA;
        }

        if (phase_ == ConnectPhase::HANDSHAKE_SEND_EP ||
            phase_ == ConnectPhase::HANDSHAKE_RECV_EP ||
            phase_ == ConnectPhase::CHANNEL_CREATING) {
            return AdvanceHandshakeAndChannel();
        }

        return hccl::HcclResult::SUCCESS;
    }


    // ============================================================
    // Internal Utilities
    // ============================================================

    // InitState：创建时初始化所有内部变量
    hccl::HcclResult HixlCsClient::InitState(const char *server_ip,
                                             uint32_t server_port,
                                             const hccl::EndPoint *end_point) {
        server_ip_ = server_ip;
        server_port_ = server_port;
        self_endpoint_ = *end_point;

        is_connected_ = false;
        socket_.reset();
        connect_started_ = false;
        phase_ = ConnectPhase::IDLE;
        send_ep_offset_ = 0;
        recv_ep_offset_ = 0;
        peer_endpoint_ = hccl::EndPoint{};
        registered_fd_ = nullptr;

        channels_.clear();
        remote_mem_info_.clear();

        return hccl::HcclResult::SUCCESS;
    }

    // CreateSelfEndpoint：创建本端 EndpointHandle
    hccl::HcclResult HixlCsClient::CreateSelfEndpoint(const hccl::EndPoint *end_point) {
        EndpointHandle handle = nullptr;
        auto ret = endpoint_store_->CreateEndpoint(*end_point, &handle);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }
        self_endpoint_handle_ = handle;
        return hccl::HcclResult::SUCCESS;
    }


    // ============================================================
    // Destroy Helpers
    // ============================================================

    // 销毁所有 endpoint/channels/socket
    hccl::HcclResult HixlCsClient::DestroyAllEndpointsAndChannels() {
        std::lock_guard<std::mutex> chnLock(chn_mutex_);

        for (auto &kv: channels_) {
            auto &info = kv.second;

            if (info.endpoint_handle && info.channel_handle) {
                endpoint_store_->DestroyChannel(info.endpoint_handle,
                                                info.channel_handle);
            }
            if (info.endpoint_handle) {
                endpoint_store_->DestroyEndpoint(info.endpoint_handle);
            }
        }
        channels_.clear();

        if (self_endpoint_handle_) {
            endpoint_store_->DestroyEndpoint(self_endpoint_handle_);
            self_endpoint_handle_ = nullptr;
        }

        if (socket_) {
            socket_->Close();
            socket_.reset();
        }

        return hccl::HcclResult::SUCCESS;
    }

    // 清理远端 mem 导入信息
    void HixlCsClient::ClearRemoteMemInfo() {
        for (auto &kv: remote_mem_info_) {
            for (auto &info: kv.second) {
                if (info.mem_tag) delete[] info.mem_tag;
            }
        }
        remote_mem_info_.clear();
    }


    // ============================================================
    // Reset
    // ============================================================

    // Reset 全部连接状态
    void HixlCsClient::ResetConnectState() {
        is_connected_ = false;
        connect_started_ = false;
        phase_ = ConnectPhase::IDLE;
        send_ep_offset_ = 0;
        recv_ep_offset_ = 0;
        peer_endpoint_ = hccl::EndPoint{};
        registered_fd_ = nullptr;
    }


    // ============================================================
    // Socket / Handshake
    // ============================================================

    // InitSocketIfNeeded：首次使用前创建 HcclSocket
    hccl::HcclResult HixlCsClient::InitSocketIfNeeded() {
        if (socket_) {
            return hccl::HcclResult::SUCCESS;
        }

        // 构造远端 IP（文本转 IP）
        hccl::HcclIpAddress remoteIp(server_ip_);

        // 用伪造的 devCtx（真实环境需要替换）
        hccl::HcclNetDevCtx netDevCtx = reinterpret_cast<hccl::HcclNetDevCtx>(0x1);

        // 创建 HcclSocket
        auto sock = std::make_shared<hccl::HcclSocket>(
            "hixl_cs_client",
            netDevCtx,
            remoteIp,
            server_port_,
            hccl::HcclSocketRole::SOCKET_ROLE_CLIENT);

        if (!sock) {
            return hccl::HcclResult::E_PARA;
        }

        // 初始化 socket（分配 NIC handle 等）
        hccl::HcclResult raRet = sock->Init();
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        socket_ = sock;
        return hccl::HcclResult::SUCCESS;
    }

    // StartRaConnect：发起非阻塞建链
    hccl::HcclResult HixlCsClient::StartRaConnect() {
        hccl::HcclResult raRet = socket_->Connect();
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }
        return hccl::HcclResult::SUCCESS;
    }

    // HandleRaConnecting：处理 connect 过程中状态
    hccl::HcclResult HixlCsClient::HandleRaConnecting(int32_t *status) {
        auto s = socket_->GetStatus();

        // 已成功连上
        if (s == hccl::HcclSocketStatus::SOCKET_OK) {
            auto ret = HandleSocketOk();
            if (ret != hccl::HcclResult::SUCCESS) {
                *status = -1;
                return ret;
            }
            *status = 0;
            return hccl::HcclResult::SUCCESS;
        }

        // 正在连接中
        if (s == hccl::HcclSocketStatus::SOCKET_CONNECTING ||
            s == hccl::HcclSocketStatus::SOCKET_INIT) {
            *status = 0;
            return hccl::HcclResult::SUCCESS;
        }

        // 状态错误
        phase_ = ConnectPhase::FAILED;
        *status = -1;
        return hccl::HcclResult::SUCCESS;
    }

    // HandleSocketOk：当 socket 连接建立成功
    hccl::HcclResult HixlCsClient::HandleSocketOk() {
        // 获取 FD
        auto fd = socket_->GetFdHandle();
        if (fd == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        // 将 fd 注册到 epoll（EventLoop）
        registered_fd_ = fd;
        auto &loop = HixlEventLoop::GetInstance();
        auto ret = loop.AddClient(fd, this);
        if (ret != hccl::HcclResult::SUCCESS) {
            phase_ = ConnectPhase::FAILED;
            return ret;
        }

        // 进入“发送 Endpoint 握手阶段”
        phase_ = ConnectPhase::HANDSHAKE_SEND_EP;
        send_ep_offset_ = 0;
        recv_ep_offset_ = 0;
        peer_endpoint_ = hccl::EndPoint{};

        return hccl::HcclResult::SUCCESS;
    }

    // 更新 status：根据状态机决定返回值
    void HixlCsClient::UpdateStatusByPhase(int32_t *status) const {
        switch (phase_) {
            case ConnectPhase::DONE:
                *status = is_connected_ ? 1 : 0;
                break;
            case ConnectPhase::FAILED:
                *status = -1;
                break;
            default:
                *status = 0;
                break;
        }
    }


    // ============================================================
    // 握手 / Channel 建立阶段
    // ============================================================

    // AdvanceHandshakeAndChannel：推进整个握手 + 建链流程
    hccl::HcclResult HixlCsClient::AdvanceHandshakeAndChannel() {
        // 1）发送本端 Endpoint
        if (phase_ == ConnectPhase::HANDSHAKE_SEND_EP) {
            auto ret = HandleHandshakeSendEp();
            if (ret != hccl::HcclResult::SUCCESS) {
                phase_ = ConnectPhase::FAILED;
                return ret;
            }
        }

        // 2）接收 server 的 Endpoint
        if (phase_ == ConnectPhase::HANDSHAKE_RECV_EP) {
            auto ret = HandleHandshakeRecvEp();
            if (ret != hccl::HcclResult::SUCCESS) {
                phase_ = ConnectPhase::FAILED;
                return ret;
            }
        }

        // 3）创建 Endpoint + Channel
        if (phase_ == ConnectPhase::CHANNEL_CREATING) {
            auto ret = HandleChannelCreating();
            if (ret != hccl::HcclResult::SUCCESS) {
                phase_ = ConnectPhase::FAILED;
                return ret;
            }
        }

        return hccl::HcclResult::SUCCESS;
    }

    // HandleHandshakeSendEp：ISend 发送本端 Endpoint
    hccl::HcclResult HixlCsClient::HandleHandshakeSendEp() {
        const char *data =
                reinterpret_cast<const char *>(&self_endpoint_) + send_ep_offset_;
        size_t remain = sizeof(self_endpoint_) - send_ep_offset_;

        // 已发送完
        if (remain == 0) {
            phase_ = ConnectPhase::HANDSHAKE_RECV_EP;
            return hccl::HcclResult::SUCCESS;
        }

        // 非阻塞发送 ISend()
        uint64_t compSize = 0;
        hccl::HcclResult raRet =
                socket_->ISend(const_cast<char *>(data), remain, compSize);
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        // 累积发送字节
        send_ep_offset_ += static_cast<size_t>(compSize);
        if (send_ep_offset_ >= sizeof(self_endpoint_)) {
            phase_ = ConnectPhase::HANDSHAKE_RECV_EP;
        }

        return hccl::HcclResult::SUCCESS;
    }

    // HandleHandshakeRecvEp：非阻塞接收服务端 Endpoint
    hccl::HcclResult HixlCsClient::HandleHandshakeRecvEp() {
        char *buf =
                reinterpret_cast<char *>(&peer_endpoint_) + recv_ep_offset_;
        size_t remain = sizeof(peer_endpoint_) - recv_ep_offset_;

        if (remain == 0) {
            phase_ = ConnectPhase::CHANNEL_CREATING;
            return hccl::HcclResult::SUCCESS;
        }

        uint64_t compSize = 0;
        hccl::HcclResult raRet =
                socket_->IRecv(buf, remain, compSize);

        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        recv_ep_offset_ += static_cast<size_t>(compSize);
        if (recv_ep_offset_ >= sizeof(peer_endpoint_)) {
            phase_ = ConnectPhase::CHANNEL_CREATING;
        }

        return hccl::HcclResult::SUCCESS;
    }

    // HandleChannelCreating：
    // ------------------------------------------------------------
    // 该函数在握手阶段最后被调用，用于“真正建立通信通道”。
    // 流程：
    // 1）使用 peer_endpoint_（从服务器接收的 Endpoint）创建一个“远端 Endpoint”
    // 2）基于“本端 self_endpoint_”和“远端 serverEp”创建一条 Channel
    // 3）将 {endpoint_handle, channel_handle, socket} 存入 channels_（key=fd）
    // 4）标记 is_connected_=true，phase_=DONE —— 整个建链完成
    // ------------------------------------------------------------
    hccl::HcclResult HixlCsClient::HandleChannelCreating() {
        // 1）创建远端 Endpoint（对应 HcommEndPointCreate）
        EndpointHandle serverEp = nullptr;
        auto ret = endpoint_store_->CreateEndpoint(peer_endpoint_, &serverEp);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }

        // 2）创建 Channel（对应 HcommChannelCreate）
        ChannelHandle ch = nullptr;
        ret = endpoint_store_->CreateChannel(serverEp, &ch);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }

        // 3）将 {endpoint, channel, socket} 注册到 channels_ 中
        auto fd = socket_->GetFdHandle();
        EndpointChannelInfo info;
        info.endpoint_handle = serverEp;
        info.channel_handle = ch;
        info.socket = socket_;
        info.is_connected = true; {
            std::lock_guard<std::mutex> lk(chn_mutex_);
            channels_[fd] = info;
        }

        // 4）建链完成
        is_connected_ = true;
        phase_ = ConnectPhase::DONE;

        return hccl::HcclResult::SUCCESS;
    }

    // FindActiveChannelUnlocked：寻找首个已连接的 channel
    EndpointChannelInfo *HixlCsClient::FindActiveChannelUnlocked() {
        for (auto &kv: channels_) {
            auto &info = kv.second;
            if (info.is_connected && info.socket) {
                return &info;
            }
        }
        return nullptr;
    }


    // ============================================================
    // GetRemoteMem 辅助函数
    // ============================================================

    hccl::HcclResult HixlCsClient::GetActiveSocketAndEndpoint(
        std::shared_ptr<hccl::HcclSocket> &socket,
        EndpointHandle &endpoint) {
        std::lock_guard<std::mutex> chnLock(chn_mutex_);
        auto *info = FindActiveChannelUnlocked();

        if (info == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        socket = info->socket;
        endpoint = info->endpoint_handle;

        if (!socket || endpoint == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        return hccl::HcclResult::SUCCESS;
    }

    // 发送 “请求远端 memdesc” 的消息
    hccl::HcclResult HixlCsClient::SendRemoteMemRequest(
        hccl::HcclSocket &socket,
        uint32_t timeout) {
        RemoteMemRequest req{};
        req.header.magic = REMOTE_MAGIC;
        req.header.version = REMOTE_VER;
        req.header.timeout = timeout;
        req.header.reserved = 0;
        req.client_ep = self_endpoint_;

        hccl::HcclResult raRet = socket.Send(&req, sizeof(req));
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        return hccl::HcclResult::SUCCESS;
    }

    // 接收共享内存数量，并分配数组
    hccl::HcclResult HixlCsClient::RecvRemoteMemCount(
        hccl::HcclSocket &socket,
        uint32_t *memCount,
        hccl::HcclMem **remote_mem_list,
        char **mem_tag_list) {
        if (!memCount || !remote_mem_list || !mem_tag_list) {
            return hccl::HcclResult::E_PARA;
        }

        *memCount = 0;

        hccl::HcclResult raRet = socket.Recv(memCount, sizeof(uint32_t));
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        if (*memCount == 0) {
            *remote_mem_list = nullptr;
            *mem_tag_list = nullptr;
            return hccl::HcclResult::SUCCESS;
        }

        return AllocateRemoteMemArrays(*memCount,
                                       remote_mem_list,
                                       mem_tag_list);
    }

    // 分配本地数组，用于保存服务端的 memdesc
    hccl::HcclResult HixlCsClient::AllocateRemoteMemArrays(
        uint32_t memCount,
        hccl::HcclMem **memList,
        char **tagBufAll) {
        try {
            *memList = new hccl::HcclMem[memCount];
            *tagBufAll = new char[memCount * TAG_MAX_LEN];
            std::memset(*tagBufAll, 0, memCount * TAG_MAX_LEN);
        } catch (...) {
            if (*memList) delete[] *memList;
            if (*tagBufAll) delete[] *tagBufAll;
            return hccl::HcclResult::E_PARA;
        }

        return hccl::HcclResult::SUCCESS;
    }

    // 接收一条 remote memdesc + tag，并执行 MemImport
    hccl::HcclResult HixlCsClient::RecvOneMemEntry(
        hccl::HcclSocket &socket,
        EndpointHandle endpoint,
        hccl::HcclMem &outMem,
        char *tagSlot,
        uint32_t tagSlotLen) {
        // 读取 descLen
        uint32_t descLen = 0;
        hccl::HcclResult raRet = socket.Recv(&descLen, sizeof(descLen));
        if (raRet != hccl::HcclResult::SUCCESS || descLen == 0) {
            return hccl::HcclResult::E_PARA;
        }

        // 读取 desc 内容
        std::vector<uint8_t> desc(descLen);
        raRet = socket.Recv(desc.data(), descLen);
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        // 读取 tag
        uint32_t tagLen = 0;
        raRet = socket.Recv(&tagLen, sizeof(tagLen));
        if (raRet != hccl::HcclResult::SUCCESS) {
            return hccl::HcclResult::E_PARA;
        }

        std::vector<char> tagTmp;
        if (tagLen > 0) {
            tagTmp.resize(tagLen + 1);
            raRet = socket.Recv(tagTmp.data(), tagLen);
            if (raRet != hccl::HcclResult::SUCCESS) {
                return hccl::HcclResult::E_PARA;
            }
            tagTmp[tagLen] = '\0';
        }

        // 执行 MemImport（desc → 可访问的内存）
        hccl::HcommBuf importedBuf{};
        auto ret = endpoint_store_->MemImport(
            endpoint, desc.data(), descLen, &importedBuf);
        if (ret != hccl::HcclResult::SUCCESS) {
            return ret;
        }

        outMem.type = hccl::HcclMemType::DEVICE;
        outMem.addr = importedBuf.addr;
        outMem.size = importedBuf.len;

        // 写入 tag
        if (tagSlot && tagSlotLen > 0) {
            if (!tagTmp.empty()) {
                std::strncpy(tagSlot, tagTmp.data(), tagSlotLen - 1);
                tagSlot[tagSlotLen - 1] = '\0';
            } else {
                tagSlot[0] = '\0';
            }
        }

        // 追加记录
        uint64_t key = reinterpret_cast<uint64_t>(importedBuf.addr);
        EndpointMemInfo info{};
        info.mem_info = outMem;
        info.mem_tag = nullptr;
        info.endpoint_handle = endpoint;
        remote_mem_info_[key].push_back(info);

        return hccl::HcclResult::SUCCESS;
    }

    // 发生错误时释放本地 buf
    void HixlCsClient::CleanupRemoteMemOnError(hccl::HcclMem *memList,
                                               char *tagBufAll) const {
        if (memList) delete[] memList;
        if (tagBufAll) delete[] tagBufAll;
    }
} // namespace hixl
